import numpy as np
import torch
import os


# trained_model=GNN.GNNClassifier(21,3,512,64,6)
trained_model = torch.load(r'.\ML\model.npt')
trained_model.eval()

A = torch.from_numpy(np.load(os.getcwd()+r'\ML\data\A.npy'))


def gen_id_2_type_dict():
    types = os.listdir(os.getcwd()+r"\ML\data\nodes")
    types = [type.split('.')[0] for type in types]
    dic = {i: name.split('_')[0] for i, name in enumerate(types)}
    return dic


id_2_type_dic = gen_id_2_type_dict()


def predict(frame):
    if len(frame) != 21:
        return None
    else:
        tensor_x = torch.from_numpy(
            np.array(frame).reshape((1, 21, 3))).to(torch.float32)
        result = trained_model(A, tensor_x)
        result = list(result.reshape(6))
        m = max(result)
        id = -1
        for i in range(len(result)):
            if result[i] == m:
                id = i
        type = id_2_type_dic[id]
        return type
